#/usr/bin/python
# -*- coding: utf-8 -*-

import parameters

import mpi4py.MPI as mpi
from solver import *
import scipy.sparse as sp
import numpy as np

from solver import Direct, GMRES, CG 

class SolvingStrategy():
    
    def __init__(self):
        pass

    def __repr__(self):
        
        return 'unimplemented SolvingStrategy'
    
    def constrainedSystem(self, A, b, C, d):
        
        if parameters.CONSTRAINTS == 'lagrange multipliers':
            
            A = sp.bmat([[A,C.T],[C,None]], "csr", dtype=parameters.FLOAT)
            b = sp.bmat([[b],[d]], "csr", dtype=parameters.FLOAT)

                
        elif parameters.CONSTRAINTS == 'substitution':
            
            pass
            
        elif parameters.CONSTRAINTS == 'penalty':
            
            pass
        
        elif parameters.CONSTRAINTS == 'elimination':
            
            pass
        
        elif self.log:
            
            self.log.debug("Constraint parameter does not match any proposal")
        
        return A, b
    
    def unConstrainedSolution(self,X,C,d,b=None, output=sp.lil_matrix):
        
        dshape = d.get_shape()
        
        Xshape = X.get_shape()
        print Xshape, dshape
        
        if parameters.CONSTRAINTS == 'lagrange multipliers':
        
            if dshape[0]==Xshape[0]:
                l = X[:,-dshape[1]:]
                X = X[:,:-dshape[1]]
            elif dshape[1]==Xshape[1]:
                l = X[-dshape[0]:,:]
                X = X[:-dshape[0],:]
            elif dshape[0]==Xshape[1]:
                l = X[-dshape[1]:,:]
                X = X[:-dshape[1],:]
            elif dshape[1]==Xshape[0]:
                l = X[:,-dshape[0]:]
                X = X[:,:-dshape[0]]           
                
            print 'lambdas = ', l.get_shape(), C.get_shape()
            
            l = C.T*l.T
            
            if not b is None:
                dshape = dshape[0]*dshape[1]
                lshape = l.get_shape()
                bshape = b.get_shape()
                if lshape[0]==bshape[0]:
                    l = l - b[:,0:-dshape]
                elif lshape[1]==bshape[1]:
                    l = l - b[0:-dshape,:]
                elif lshape[0]==bshape[1]:
                    l = l - b[0:-dshape,:].T
                elif lshape[1]==bshape[0]:
                    l = l - b[:,0:-dshape].T
            
        return output(X), output(l)

    def solve(self):
        '''
            
        '''
        
        return None, None
    
class FullModel(SolvingStrategy):
    
    def __init__(self, solver=Direct, log=None):
        self.log = log
        self.solver = solver()
    
    def __repr__(self):
        
        return 'full solution strategy'
    
    def computeError(self):
        pass
    
    def computeOperators(self, domain, lhslist, rhslist, disp=1, load=1, meshId=None, error=True):
        # operators:
        lhs, rhs = domain.computeOperators(lhslist, rhslist, disp, load, meshId)
        
        if error:
            try:
                res = np.sqrt(np.asscalar(rhs['R'].T*rhs['R'])/np.asscalar(rhs['Fext'].T*rhs['Fext']))
                return res
            except:
                print "warning: error measure failed"
                return 0
        else:
            return 0
           
    def solve(self, domain, lhslist, rhslist, meshId=None):
        
        # constraints:
        constraints = domain.getConstraints(full=True)
        
        # linear system solver:
        A = domain.getLhs(lhslist[0])
        b = domain.getRhs(rhslist[0])

        # Constrained system:
        if len(constraints.Cd)>1: # multiple load cases
            pass
#        for i in xrange(len(Cf)):
#            b = sp.bmat([[b, Cf[i].T*f[i]]])
#        
#        for j in xrange(len(Cd)):   
#            Aj, bj = self.constrainedSystem(A, b, Cd, d)
        else: # One load case  
        
            A, b = self.constrainedSystem(A, b, constraints.Cd[0], constraints.d[0])
            
        # solve system
        X = self.solver.solve(A,b)
        
        X, lam = self.unConstrainedSolution(X, constraints.Cd[0], constraints.d[0], b)
        
        return X, lam
    
class DomainDecomposed(SolvingStrategy):
    
    def __init__(self, gsolver=Direct, lsolver=Direct, log=None):
        self.log = log
        self.globalSolver = gsolver()
        self.localSolver = lsolver()
    
    def __repr__(self):
        
        return 'domain decomposition solution strategy'
    
    def computeOperators(self, domain, lhslist, rhslist, disp=1, load=1, meshId=None):
        # operators:
        domain.computeOperators(lhslist, rhslist, disp, load, meshId, updateDofs)
        
    
    def solve(self, domain, lhslist, rhslist,  disp=1, load=1, distributed=True):
        
        ## Global step:
        self.globalStep(domain, lhslist, rhslist)
        
        # Local step 
        domain.localize(self.localStep, lhslist, rhslist, self.localSolver)
        
    
    def globalStep(self, domain, lhslist, rhslist):
        '''
            Global step of the domain decomposition solver.
        '''
        
        if distributed:    # solving without assembly of the global problem: 
            Xg = self.globalSolver(domain,lhslist,rhslist,)
        else:              # solving the assembled global problem:
            # constraint on global solution:
            constraints = domain.getConstraints(full=False)
            # assembly of Schur complements:
            A = domain.getLhs(lhslist,full=False)
            b = domain.getRhs(rhslist,full=False)
            # constrained system:
            A, b = self.constrainedSystem(A, b, constraints.Cd[0], constraints.d[0])
            # solve system
            Xg = self.solver.solve(A,b)
            Xg, lam = self.unConstrainedSolution(Xg,constraints.d[0])
            
            # distribute solution over subdomains
            domain.setGlobalDofs(Xg)
        
        
    @staticmethod
    def localStep(subdomain, *args): 
        '''
            Local step of the domain decomposition solver.
        '''
        
        pass
        
          
    
class ReducedModel(SolvingStrategy):

    def __init__(self):
        pass
    
    def solve(self):
        pass
    
    
    
    
    
    